{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "6NMSi4HFFtJk"
            },
            "source": [
                "# 1. Preparations\n",
                "\n",
                "Before start, make sure that you choose\n",
                "* Runtime Type = Python 3\n",
                "* Hardware Accelerator = GPU\n",
                "\n",
                "in the **Runtime** menu -> **Change runtime type**\n",
                "\n",
                "Then, we set up the envrironment.\n",
                "\n",
                "* install bmf\n",
                "* install requirement python module\n",
                "* install wurlitzer and load_ext for outputing logs of c extension module"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {
                "colab": {
                    "base_uri": "https://localhost:8080/"
                },
                "executionInfo": {
                    "elapsed": 24913,
                    "status": "ok",
                    "timestamp": 1685674075625,
                    "user": {
                        "displayName": "hu heng",
                        "userId": "04066773367750540207"
                    },
                    "user_tz": -480
                },
                "id": "TWok65-nFlLn",
                "outputId": "786a506a-aea7-45e9-f951-d29dcdf2eaf8"
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Downloading...\n",
                        "From: https://drive.google.com/uc?id=1B2iUlDEUuled09IFYE3C2mD1Rx1LEmd6\n",
                        "To: /content/bmf-3.0.0-cp310-none-linux_x86_64.whl\n",
                        "100% 98.3M/98.3M [00:03<00:00, 31.2MB/s]\n",
                        "Downloading...\n",
                        "From: https://drive.google.com/uc?id=1vvUssyc8GC8SPzVdYwRwKc2OvyKGY-Px\n",
                        "To: /content/bbb_360_20s.mp4\n",
                        "100% 1.13M/1.13M [00:00<00:00, 151MB/s]\n",
                        "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
                        "Processing ./bmf-3.0.0-cp310-none-linux_x86_64.whl\n",
                        "Collecting basicsr==1.4.2\n",
                        "  Downloading basicsr-1.4.2.tar.gz (172 kB)\n",
                        "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m172.5/172.5 kB\u001b[0m \u001b[31m14.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
                        "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
                        "Collecting realesrgan\n",
                        "  Downloading realesrgan-0.3.0-py3-none-any.whl (26 kB)\n",
                        "Collecting addict (from basicsr==1.4.2)\n",
                        "  Downloading addict-2.4.0-py3-none-any.whl (3.8 kB)\n",
                        "Requirement already satisfied: future in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (0.18.3)\n",
                        "Collecting lmdb (from basicsr==1.4.2)\n",
                        "  Downloading lmdb-1.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (299 kB)\n",
                        "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m299.2/299.2 kB\u001b[0m \u001b[31m28.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
                        "\u001b[?25hRequirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (1.22.4)\n",
                        "Requirement already satisfied: opencv-python in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (4.7.0.72)\n",
                        "Requirement already satisfied: Pillow in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (8.4.0)\n",
                        "Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (6.0)\n",
                        "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (2.27.1)\n",
                        "Requirement already satisfied: scikit-image in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (0.19.3)\n",
                        "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (1.10.1)\n",
                        "Collecting tb-nightly (from basicsr==1.4.2)\n",
                        "  Downloading tb_nightly-2.14.0a20230601-py3-none-any.whl (5.6 MB)\n",
                        "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.6/5.6 MB\u001b[0m \u001b[31m20.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
                        "\u001b[?25hRequirement already satisfied: torch>=1.7 in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (2.0.1+cu118)\n",
                        "Requirement already satisfied: torchvision in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (0.15.2+cu118)\n",
                        "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from basicsr==1.4.2) (4.65.0)\n",
                        "Collecting yapf (from basicsr==1.4.2)\n",
                        "  Downloading yapf-0.33.0-py2.py3-none-any.whl (200 kB)\n",
                        "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m200.9/200.9 kB\u001b[0m \u001b[31m25.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
                        "\u001b[?25hCollecting facexlib>=0.2.5 (from realesrgan)\n",
                        "  Downloading facexlib-0.3.0-py3-none-any.whl (59 kB)\n",
                        "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m59.6/59.6 kB\u001b[0m \u001b[31m8.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
                        "\u001b[?25hCollecting gfpgan>=1.3.5 (from realesrgan)\n",
                        "  Downloading gfpgan-1.3.8-py3-none-any.whl (52 kB)\n",
                        "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m52.2/52.2 kB\u001b[0m \u001b[31m7.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
                        "\u001b[?25hCollecting filterpy (from facexlib>=0.2.5->realesrgan)\n",
                        "  Downloading filterpy-1.4.5.zip (177 kB)\n",
                        "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m178.0/178.0 kB\u001b[0m \u001b[31m24.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
                        "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
                        "Requirement already satisfied: numba in /usr/local/lib/python3.10/dist-packages (from facexlib>=0.2.5->realesrgan) (0.56.4)\n",
                        "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=1.7->basicsr==1.4.2) (3.12.0)\n",
                        "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch>=1.7->basicsr==1.4.2) (4.5.0)\n",
                        "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.7->basicsr==1.4.2) (1.11.1)\n",
                        "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.7->basicsr==1.4.2) (3.1)\n",
                        "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.7->basicsr==1.4.2) (3.1.2)\n",
                        "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.7->basicsr==1.4.2) (2.0.0)\n",
                        "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.7->basicsr==1.4.2) (3.25.2)\n",
                        "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.7->basicsr==1.4.2) (16.0.5)\n",
                        "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->basicsr==1.4.2) (1.26.15)\n",
                        "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->basicsr==1.4.2) (2022.12.7)\n",
                        "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->basicsr==1.4.2) (2.0.12)\n",
                        "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->basicsr==1.4.2) (3.4)\n",
                        "Requirement already satisfied: imageio>=2.4.1 in /usr/local/lib/python3.10/dist-packages (from scikit-image->basicsr==1.4.2) (2.25.1)\n",
                        "Requirement already satisfied: tifffile>=2019.7.26 in /usr/local/lib/python3.10/dist-packages (from scikit-image->basicsr==1.4.2) (2023.4.12)\n",
                        "Requirement already satisfied: PyWavelets>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from scikit-image->basicsr==1.4.2) (1.4.1)\n",
                        "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from scikit-image->basicsr==1.4.2) (23.1)\n",
                        "Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.10/dist-packages (from tb-nightly->basicsr==1.4.2) (1.4.0)\n",
                        "Requirement already satisfied: grpcio>=1.48.2 in /usr/local/lib/python3.10/dist-packages (from tb-nightly->basicsr==1.4.2) (1.54.0)\n",
                        "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.10/dist-packages (from tb-nightly->basicsr==1.4.2) (2.17.3)\n",
                        "Requirement already satisfied: google-auth-oauthlib<1.1,>=0.5 in /usr/local/lib/python3.10/dist-packages (from tb-nightly->basicsr==1.4.2) (1.0.0)\n",
                        "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.10/dist-packages (from tb-nightly->basicsr==1.4.2) (3.4.3)\n",
                        "Requirement already satisfied: protobuf>=3.19.6 in /usr/local/lib/python3.10/dist-packages (from tb-nightly->basicsr==1.4.2) (3.20.3)\n",
                        "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.10/dist-packages (from tb-nightly->basicsr==1.4.2) (67.7.2)\n",
                        "Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from tb-nightly->basicsr==1.4.2) (0.7.0)\n",
                        "Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from tb-nightly->basicsr==1.4.2) (2.3.0)\n",
                        "Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.10/dist-packages (from tb-nightly->basicsr==1.4.2) (0.40.0)\n",
                        "Requirement already satisfied: tomli>=2.0.1 in /usr/local/lib/python3.10/dist-packages (from yapf->basicsr==1.4.2) (2.0.1)\n",
                        "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tb-nightly->basicsr==1.4.2) (5.3.0)\n",
                        "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tb-nightly->basicsr==1.4.2) (0.3.0)\n",
                        "Requirement already satisfied: six>=1.9.0 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tb-nightly->basicsr==1.4.2) (1.16.0)\n",
                        "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tb-nightly->basicsr==1.4.2) (4.9)\n",
                        "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from google-auth-oauthlib<1.1,>=0.5->tb-nightly->basicsr==1.4.2) (1.3.1)\n",
                        "Requirement already satisfied: MarkupSafe>=2.1.1 in /usr/local/lib/python3.10/dist-packages (from werkzeug>=1.0.1->tb-nightly->basicsr==1.4.2) (2.1.2)\n",
                        "Requirement already satisfied: matplotlib in /usr/local/lib/python3.10/dist-packages (from filterpy->facexlib>=0.2.5->realesrgan) (3.7.1)\n",
                        "Requirement already satisfied: llvmlite<0.40,>=0.39.0dev0 in /usr/local/lib/python3.10/dist-packages (from numba->facexlib>=0.2.5->realesrgan) (0.39.1)\n",
                        "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.7->basicsr==1.4.2) (1.3.0)\n",
                        "Requirement already satisfied: pyasn1<0.6.0,>=0.4.6 in /usr/local/lib/python3.10/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tb-nightly->basicsr==1.4.2) (0.5.0)\n",
                        "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.10/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<1.1,>=0.5->tb-nightly->basicsr==1.4.2) (3.2.2)\n",
                        "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->filterpy->facexlib>=0.2.5->realesrgan) (1.0.7)\n",
                        "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib->filterpy->facexlib>=0.2.5->realesrgan) (0.11.0)\n",
                        "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->filterpy->facexlib>=0.2.5->realesrgan) (4.39.3)\n",
                        "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->filterpy->facexlib>=0.2.5->realesrgan) (1.4.4)\n",
                        "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->filterpy->facexlib>=0.2.5->realesrgan) (3.0.9)\n",
                        "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib->filterpy->facexlib>=0.2.5->realesrgan) (2.8.2)\n",
                        "Building wheels for collected packages: basicsr, filterpy\n",
                        "  Building wheel for basicsr (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
                        "  Created wheel for basicsr: filename=basicsr-1.4.2-py3-none-any.whl size=214823 sha256=4b9bb07b261f13514e341cb70222852e9e0548ba41b739c67d88789d14be2fe5\n",
                        "  Stored in directory: /root/.cache/pip/wheels/38/83/99/2d8437cc652a01af27df5ff037a4075e95b52d67705c5f30ca\n",
                        "  Building wheel for filterpy (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
                        "  Created wheel for filterpy: filename=filterpy-1.4.5-py3-none-any.whl size=110459 sha256=5958cff801e80d3021255cec838168bf2e36012b25f888b8e43a8ee01f8f5eec\n",
                        "  Stored in directory: /root/.cache/pip/wheels/0f/0c/ea/218f266af4ad626897562199fbbcba521b8497303200186102\n",
                        "Successfully built basicsr filterpy\n",
                        "Installing collected packages: lmdb, addict, yapf, bmf, filterpy, tb-nightly, facexlib, basicsr, gfpgan, realesrgan\n",
                        "Successfully installed addict-2.4.0 basicsr-1.4.2 bmf-3.0.0 facexlib-0.3.0 filterpy-1.4.5 gfpgan-1.3.8 lmdb-1.4.1 realesrgan-0.3.0 tb-nightly-2.14.0a20230601 yapf-0.33.0\n",
                        "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
                        "Collecting wurlitzer\n",
                        "  Downloading wurlitzer-3.0.3-py3-none-any.whl (7.3 kB)\n",
                        "Installing collected packages: wurlitzer\n",
                        "Successfully installed wurlitzer-3.0.3\n"
                    ]
                }
            ],
            "source": [
                "!gdown --fuzzy https://drive.google.com/file/d/1vvUssyc8GC8SPzVdYwRwKc2OvyKGY-Px/view?usp=sharing\n",
                "!pip install BabitMF basicsr==1.4.2 realesrgan torchvision==0.15.2\n",
                "!pip install wurlitzer\n",
                "%load_ext wurlitzer"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "VqNE_a5PF1NV"
            },
            "source": [
                "# 2 source code"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "-GYpVG0fx0FO"
            },
            "source": [
                "## **enhance_module**\n",
                "\n",
                "We implemente the inference process of this project [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) as a BMF module\n",
                "\n",
                "import and some util functions:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {
                "colab": {
                    "base_uri": "https://localhost:8080/"
                },
                "executionInfo": {
                    "elapsed": 7875,
                    "status": "ok",
                    "timestamp": 1685674092169,
                    "user": {
                        "displayName": "hu heng",
                        "userId": "04066773367750540207"
                    },
                    "user_tz": -480
                },
                "id": "v-DNHit0F31S",
                "outputId": "7c2fcf93-1e7a-4787-ff76-b3115738f9a3"
            },
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional_tensor.py:5: UserWarning: The torchvision.transforms.functional_tensor module is deprecated in 0.15 and will be **removed in 0.17**. Please don't rely on it. You probably just need to use APIs in torchvision.transforms.functional or in torchvision.transforms.v2.functional.\n",
                        "  warnings.warn(\n"
                    ]
                }
            ],
            "source": [
                "#!/usr/bin/env python\n",
                "# -*- coding: utf-8 -*-\n",
                "\n",
                "# This code is based on the implementation from https://github.com/xinntao/Real-ESRGAN\n",
                "\n",
                "from basicsr.archs.rrdbnet_arch import RRDBNet\n",
                "from basicsr.utils.download_util import load_file_from_url\n",
                "\n",
                "from realesrgan import RealESRGANer\n",
                "from realesrgan.archs.srvgg_arch import SRVGGNetCompact\n",
                "\n",
                "from bmf import Module, Log, Timestamp, ProcessResult, LogLevel, Packet, VideoFrame\n",
                "from bmf.lib._bmf.sdk import ffmpeg\n",
                "\n",
                "from bmf import hmp as mp\n",
                "import numpy as np\n",
                "\n",
                "import os\n",
                "\n",
                "\n",
                "def load_model():\n",
                "    model = SRVGGNetCompact(\n",
                "        num_in_ch=3,\n",
                "        num_out_ch=3,\n",
                "        num_feat=64,\n",
                "        num_conv=16,\n",
                "        upscale=4,\n",
                "        act_type=\"prelu\",\n",
                "    )\n",
                "    netscale = 4\n",
                "    file_url = [\n",
                "        \"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth\"\n",
                "    ]\n",
                "    return model, netscale, file_url\n",
                "\n",
                "\n",
                "def prepare_model(model_name, file_url):\n",
                "    model_path = os.path.join(\"weights\", model_name + \".pth\")\n",
                "    if not os.path.isfile(model_path):\n",
                "        ROOT_DIR = \"/content\"\n",
                "        for url in file_url:\n",
                "            # model_path will be updated\n",
                "            model_path = load_file_from_url(\n",
                "                url=url,\n",
                "                model_dir=os.path.join(ROOT_DIR, \"weights\"),\n",
                "                progress=True,\n",
                "                file_name=None,\n",
                "            )\n",
                "    return model_path"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "8G6JKUzC12V7"
            },
            "source": [
                "This module retrieves BMF VideoFrames from the upstream module, then uses the functions provided by the BMF framework to convert VideoFrame to RGB24 format, and then converts it to a numpy array. The numpy array is then sent to the inference object for super-resolution processing. The inference result is a RGB24 frame, which is also a numpy array. The result is then converted back to BMF VideoFrame and sent to the downstream module. This is a typical BMF module processing flow."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {
                "executionInfo": {
                    "elapsed": 19,
                    "status": "ok",
                    "timestamp": 1685674092169,
                    "user": {
                        "displayName": "hu heng",
                        "userId": "04066773367750540207"
                    },
                    "user_tz": -480
                },
                "id": "SOWHWx7e15nT"
            },
            "outputs": [],
            "source": [
                "class EnhanceModule(Module):\n",
                "    def __init__(self, node=None, option=None):\n",
                "        self._node = node\n",
                "        if not option:\n",
                "            Log.log_node(LogLevel.ERROR, self._node, \"no option\")\n",
                "            return\n",
                "\n",
                "        tile = option.get(\"tile\", 0)\n",
                "        tile_pad = option.get(\"tile_pad\", 10)\n",
                "        pre_pad = option.get(\"pre_pad\", 10)\n",
                "        fp32 = option.get(\"fp32\", False)\n",
                "        gpu_id = option.get(\"gpu_id\", 0)\n",
                "\n",
                "        self.output_scale = option.get(\"output_scale\", None)\n",
                "\n",
                "        model, netscale, file_url = load_model()\n",
                "        model_name = \"realesr-animevideov3\"  # x4 VGG-style model (XS size)\n",
                "        model_path = prepare_model(model_name, file_url)\n",
                "\n",
                "        self.upsampler = RealESRGANer(\n",
                "            scale=netscale,\n",
                "            model_path=model_path,\n",
                "            dni_weight=None,\n",
                "            model=model,\n",
                "            tile=tile,\n",
                "            tile_pad=tile_pad,\n",
                "            pre_pad=pre_pad,\n",
                "            half=not fp32,\n",
                "            gpu_id=gpu_id,\n",
                "        )\n",
                "        self.process_num = 0\n",
                "\n",
                "    def process(self, task):\n",
                "        output_queue = task.get_outputs().get(0, None)\n",
                "        input_queue = task.get_inputs().get(0, None)\n",
                "\n",
                "        while not input_queue.empty():\n",
                "            pkt = input_queue.get()\n",
                "            # process EOS\n",
                "            if pkt.timestamp == Timestamp.EOF:\n",
                "                Log.log_node(LogLevel.INFO, task.get_node(), \"Receive EOF\")\n",
                "                if output_queue is not None:\n",
                "                    output_queue.put(Packet.generate_eof_packet())\n",
                "                task.timestamp = Timestamp.DONE\n",
                "                return ProcessResult.OK\n",
                "\n",
                "            video_frame = pkt.get(VideoFrame)\n",
                "            # use ffmpeg\n",
                "            frame = ffmpeg.reformat(video_frame, \"rgb24\").frame().plane(0).numpy()\n",
                "\n",
                "            output, _ = self.upsampler.enhance(frame, self.output_scale)\n",
                "            Log.log_node(\n",
                "                LogLevel.DEBUG,\n",
                "                self._node,\n",
                "                \"enhance output shape: \",\n",
                "                output.shape,\n",
                "                \" flags: \",\n",
                "                output.flags,\n",
                "            )\n",
                "            self.process_num += 1\n",
                "            if self.process_num % 10 == 0:\n",
                "                Log.log_node(LogLevel.INFO, self._node, \"process num: \", self.process_num)\n",
                "            output = np.ascontiguousarray(output)\n",
                "            rgbformat = mp.PixelInfo(mp.kPF_RGB24)\n",
                "            image = mp.Frame(mp.from_numpy(output), rgbformat)\n",
                "\n",
                "            output_frame = VideoFrame(image)\n",
                "            Log.log_node(LogLevel.DEBUG, self._node, \"output video frame\")\n",
                "\n",
                "            output_frame.pts = video_frame.pts\n",
                "            output_frame.time_base = video_frame.time_base\n",
                "            output_pkt = Packet(output_frame)\n",
                "            output_pkt.timestamp = pkt.timestamp\n",
                "            if output_queue is not None:\n",
                "                output_queue.put(output_pkt)\n",
                "\n",
                "        return ProcessResult.OK"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "aBzWmrOc3vF9"
            },
            "source": [
                "## **compose_module**\n",
                "We implement another module which named compose_module to overlay and display the super-resolution results and the original video. \n",
                "\n",
                "This module processes two VideoFrames from the upstream module at the same time, which are the original video frame and the super-resolution frame. Then, it uses OpenCV to combine them into one frame and outputs it to the downstream module. The BMF framework provides an input manager called 'framesync', which can synchronize multiple inputs to a module. It only needs to be specified when creating the module."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "metadata": {
                "executionInfo": {
                    "elapsed": 19,
                    "status": "ok",
                    "timestamp": 1685674092170,
                    "user": {
                        "displayName": "hu heng",
                        "userId": "04066773367750540207"
                    },
                    "user_tz": -480
                },
                "id": "C9veI_dVGMMV"
            },
            "outputs": [],
            "source": [
                "class Compositor:\n",
                "    def __init__(self):\n",
                "        self.slide_step = 5\n",
                "        self.line_width = 10\n",
                "        self.frame_width = 0\n",
                "        self.frame_num = 0\n",
                "\n",
                "    def merge(self, img_left, img_right, x_coord):\n",
                "        img_out = np.copy(img_left)\n",
                "        img_out[:, x_coord:, :] = img_right[:, x_coord:, :]\n",
                "\n",
                "        # white line\n",
                "        img_out[:, x_coord : x_coord + self.line_width, :] = 255\n",
                "\n",
                "        return img_out\n",
                "\n",
                "    def compose(self, img_left, img_right):\n",
                "        if self.frame_width == 0:\n",
                "            self.frame_width = img_left.shape[1]\n",
                "\n",
                "        remainder = (self.frame_num * self.slide_step) % (2 * self.frame_width)\n",
                "        x_coord = min(remainder, 2 * self.frame_width - remainder)\n",
                "\n",
                "        out = self.merge(img_left, img_right, x_coord)\n",
                "\n",
                "        self.frame_num += 1\n",
                "\n",
                "        return out\n",
                "\n",
                "\n",
                "class CompositionModule(Module):\n",
                "    def __init__(self, node=None, option=None):\n",
                "        self._node = node\n",
                "        self.compositor = Compositor()\n",
                "\n",
                "    def process(self, task):\n",
                "        output_queue = task.get_outputs().get(0, None)\n",
                "\n",
                "        current_pts = None\n",
                "        current_timestamp = 0\n",
                "        current_timebase = None\n",
                "\n",
                "        # collect frames\n",
                "        frame_list = []\n",
                "        for (index, input_packets) in task.get_inputs().items():\n",
                "            Log.log_node(\n",
                "                LogLevel.DEBUG,\n",
                "                self._node,\n",
                "                \"process input index: \",\n",
                "                index,\n",
                "            )\n",
                "\n",
                "            num = 0\n",
                "            while not input_packets.empty():\n",
                "                pkt = input_packets.get()\n",
                "                num += 1\n",
                "                # process EOS\n",
                "                if pkt.timestamp == Timestamp.EOF:\n",
                "                    Log.log_node(\n",
                "                        LogLevel.INFO, task.get_node(), \"Receive EOF.................\"\n",
                "                    )\n",
                "                    if output_queue is not None:\n",
                "                        output_queue.put(Packet.generate_eof_packet())\n",
                "                    task.timestamp = Timestamp.DONE\n",
                "                    return ProcessResult.OK\n",
                "\n",
                "                frame = pkt.get(VideoFrame)\n",
                "\n",
                "                # record\n",
                "                if current_pts is None:\n",
                "                    current_pts = frame.pts\n",
                "                    current_timebase = frame.time_base\n",
                "                    current_timestamp = pkt.timestamp\n",
                "\n",
                "                rgb_frame = ffmpeg.reformat(frame, \"rgb24\").frame().plane(0).numpy()\n",
                "\n",
                "                Log.log_node(\n",
                "                    LogLevel.DEBUG,\n",
                "                    self._node,\n",
                "                    \"index:\",\n",
                "                    index,\n",
                "                    \"get frame pts: \",\n",
                "                    frame.pts,\n",
                "                )\n",
                "                frame_list.append(rgb_frame)\n",
                "\n",
                "            Log.log_node(\n",
                "                LogLevel.DEBUG, self._node, \"index:\", index, \"get frame num:\", num\n",
                "            )\n",
                "\n",
                "        if len(frame_list) != 2:\n",
                "            Log.log_node(\n",
                "                LogLevel.WARNING,\n",
                "                self._node,\n",
                "                \"frame list num incorrect, check framesync input\",\n",
                "            )\n",
                "            return ProcessResult.OK\n",
                "\n",
                "        out = self.compositor.compose(frame_list[0], frame_list[1])\n",
                "\n",
                "        rgbformat = mp.PixelInfo(mp.kPF_RGB24)\n",
                "        image = mp.Frame(mp.from_numpy(out), rgbformat)\n",
                "        output_frame = VideoFrame(image)\n",
                "        output_frame.pts = current_pts\n",
                "        output_frame.time_base = current_timebase\n",
                "        output_pkt = Packet(output_frame)\n",
                "        output_pkt.timestamp = current_timestamp\n",
                "\n",
                "        # output index 0 for test\n",
                "        if output_queue is not None:\n",
                "            output_queue.put(output_pkt)\n",
                "\n",
                "        return ProcessResult.OK"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "fYmlaQWF7DSd"
            },
            "source": [
                "## **main program**\n",
                "\n",
                "We create module instances and connect those modules into a pipeline through a main program, run it, and finally obtain the results.\n",
                "\n",
                "![enhance_graph.png]()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "metadata": {
                "executionInfo": {
                    "elapsed": 18,
                    "status": "ok",
                    "timestamp": 1685674092170,
                    "user": {
                        "displayName": "hu heng",
                        "userId": "04066773367750540207"
                    },
                    "user_tz": -480
                },
                "id": "TapNGQFYHy-v"
            },
            "outputs": [],
            "source": [
                "Log.set_log_level(LogLevel.INFO)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "metadata": {
                "colab": {
                    "base_uri": "https://localhost:8080/"
                },
                "executionInfo": {
                    "elapsed": 164635,
                    "status": "ok",
                    "timestamp": 1685674256787,
                    "user": {
                        "displayName": "hu heng",
                        "userId": "04066773367750540207"
                    },
                    "user_tz": -480
                },
                "id": "F7Qxpx6tVG_d",
                "outputId": "109bb362-5e17-445f-da55-246fb2f52125"
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "{\n",
                        "    \"input_streams\": [],\n",
                        "    \"output_streams\": [],\n",
                        "    \"nodes\": [\n",
                        "        {\n",
                        "            \"module_info\": {\n",
                        "                \"name\": \"c_ffmpeg_decoder\",\n",
                        "                \"type\": \"\",\n",
                        "                \"path\": \"\",\n",
                        "                \"entry\": \"\"\n",
                        "            },\n",
                        "            \"meta_info\": {\n",
                        "                \"premodule_id\": -1,\n",
                        "                \"callback_binding\": []\n",
                        "            },\n",
                        "            \"option\": {\n",
                        "                \"input_path\": \"bbb_360_20s.mp4\"\n",
                        "            },\n",
                        "            \"input_streams\": [],\n",
                        "            \"output_streams\": [\n",
                        "                {\n",
                        "                    \"identifier\": \"video:c_ffmpeg_decoder_0_1\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                },\n",
                        "                {\n",
                        "                    \"identifier\": \"audio:c_ffmpeg_decoder_0_2\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                }\n",
                        "            ],\n",
                        "            \"input_manager\": \"immediate\",\n",
                        "            \"scheduler\": 0,\n",
                        "            \"alias\": \"\",\n",
                        "            \"id\": 0\n",
                        "        },\n",
                        "        {\n",
                        "            \"module_info\": {\n",
                        "                \"name\": \"EnhanceModule\",\n",
                        "                \"type\": \"\",\n",
                        "                \"path\": \"\",\n",
                        "                \"entry\": \"__main__.EnhanceModule\"\n",
                        "            },\n",
                        "            \"meta_info\": {\n",
                        "                \"premodule_id\": -1,\n",
                        "                \"callback_binding\": []\n",
                        "            },\n",
                        "            \"option\": {\n",
                        "                \"fp32\": true,\n",
                        "                \"output_scale\": 2\n",
                        "            },\n",
                        "            \"input_streams\": [\n",
                        "                {\n",
                        "                    \"identifier\": \"c_ffmpeg_decoder_0_1\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                }\n",
                        "            ],\n",
                        "            \"output_streams\": [\n",
                        "                {\n",
                        "                    \"identifier\": \"EnhanceModule_1_0\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                }\n",
                        "            ],\n",
                        "            \"input_manager\": \"immediate\",\n",
                        "            \"scheduler\": 0,\n",
                        "            \"alias\": \"\",\n",
                        "            \"id\": 1\n",
                        "        },\n",
                        "        {\n",
                        "            \"module_info\": {\n",
                        "                \"name\": \"c_ffmpeg_filter\",\n",
                        "                \"type\": \"\",\n",
                        "                \"path\": \"\",\n",
                        "                \"entry\": \"\"\n",
                        "            },\n",
                        "            \"meta_info\": {\n",
                        "                \"premodule_id\": -1,\n",
                        "                \"callback_binding\": []\n",
                        "            },\n",
                        "            \"option\": {\n",
                        "                \"name\": \"scale\",\n",
                        "                \"para\": \"2*in_w:2*in_h\"\n",
                        "            },\n",
                        "            \"input_streams\": [\n",
                        "                {\n",
                        "                    \"identifier\": \"c_ffmpeg_decoder_0_1\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                }\n",
                        "            ],\n",
                        "            \"output_streams\": [\n",
                        "                {\n",
                        "                    \"identifier\": \"c_ffmpeg_filter_2_0\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                }\n",
                        "            ],\n",
                        "            \"input_manager\": \"immediate\",\n",
                        "            \"scheduler\": 0,\n",
                        "            \"alias\": \"\",\n",
                        "            \"id\": 2\n",
                        "        },\n",
                        "        {\n",
                        "            \"module_info\": {\n",
                        "                \"name\": \"CompositionModule\",\n",
                        "                \"type\": \"\",\n",
                        "                \"path\": \"\",\n",
                        "                \"entry\": \"__main__.CompositionModule\"\n",
                        "            },\n",
                        "            \"meta_info\": {\n",
                        "                \"premodule_id\": -1,\n",
                        "                \"callback_binding\": []\n",
                        "            },\n",
                        "            \"option\": {},\n",
                        "            \"input_streams\": [\n",
                        "                {\n",
                        "                    \"identifier\": \"EnhanceModule_1_0\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                },\n",
                        "                {\n",
                        "                    \"identifier\": \"c_ffmpeg_filter_2_0\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                }\n",
                        "            ],\n",
                        "            \"output_streams\": [\n",
                        "                {\n",
                        "                    \"identifier\": \"CompositionModule_3_0\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                }\n",
                        "            ],\n",
                        "            \"input_manager\": \"framesync\",\n",
                        "            \"scheduler\": 0,\n",
                        "            \"alias\": \"\",\n",
                        "            \"id\": 3\n",
                        "        },\n",
                        "        {\n",
                        "            \"module_info\": {\n",
                        "                \"name\": \"c_ffmpeg_encoder\",\n",
                        "                \"type\": \"\",\n",
                        "                \"path\": \"\",\n",
                        "                \"entry\": \"\"\n",
                        "            },\n",
                        "            \"meta_info\": {\n",
                        "                \"premodule_id\": -1,\n",
                        "                \"callback_binding\": []\n",
                        "            },\n",
                        "            \"option\": {\n",
                        "                \"video_params\": {\n",
                        "                    \"g\": \"50\",\n",
                        "                    \"preset\": \"veryfast\",\n",
                        "                    \"bf\": 0,\n",
                        "                    \"vsync\": \"vfr\",\n",
                        "                    \"max_fr\": 30\n",
                        "                },\n",
                        "                \"audio_params\": {\n",
                        "                    \"sample_rate\": 44100,\n",
                        "                    \"codec\": \"aac\"\n",
                        "                },\n",
                        "                \"loglevel\": \"info\",\n",
                        "                \"output_path\": \"output.mp4\"\n",
                        "            },\n",
                        "            \"input_streams\": [\n",
                        "                {\n",
                        "                    \"identifier\": \"EnhanceModule_1_0\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                },\n",
                        "                {\n",
                        "                    \"identifier\": \"c_ffmpeg_decoder_0_2\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                }\n",
                        "            ],\n",
                        "            \"output_streams\": [],\n",
                        "            \"input_manager\": \"immediate\",\n",
                        "            \"scheduler\": 1,\n",
                        "            \"alias\": \"\",\n",
                        "            \"id\": 4\n",
                        "        },\n",
                        "        {\n",
                        "            \"module_info\": {\n",
                        "                \"name\": \"c_ffmpeg_encoder\",\n",
                        "                \"type\": \"\",\n",
                        "                \"path\": \"\",\n",
                        "                \"entry\": \"\"\n",
                        "            },\n",
                        "            \"meta_info\": {\n",
                        "                \"premodule_id\": -1,\n",
                        "                \"callback_binding\": []\n",
                        "            },\n",
                        "            \"option\": {\n",
                        "                \"video_params\": {\n",
                        "                    \"g\": \"50\",\n",
                        "                    \"preset\": \"veryfast\",\n",
                        "                    \"bf\": 0,\n",
                        "                    \"vsync\": \"vfr\",\n",
                        "                    \"max_fr\": 30\n",
                        "                },\n",
                        "                \"audio_params\": {\n",
                        "                    \"sample_rate\": 44100,\n",
                        "                    \"codec\": \"aac\"\n",
                        "                },\n",
                        "                \"loglevel\": \"info\",\n",
                        "                \"output_path\": \"output_compose.mp4\"\n",
                        "            },\n",
                        "            \"input_streams\": [\n",
                        "                {\n",
                        "                    \"identifier\": \"CompositionModule_3_0\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                },\n",
                        "                {\n",
                        "                    \"identifier\": \"c_ffmpeg_decoder_0_2\",\n",
                        "                    \"stream_alias\": \"\"\n",
                        "                }\n",
                        "            ],\n",
                        "            \"output_streams\": [],\n",
                        "            \"input_manager\": \"immediate\",\n",
                        "            \"scheduler\": 1,\n",
                        "            \"alias\": \"\",\n",
                        "            \"id\": 5\n",
                        "        }\n",
                        "    ],\n",
                        "    \"option\": {},\n",
                        "    \"mode\": \"Normal\"\n",
                        "}\n",
                        "[2023-06-02 02:48:11.741] [info] debug module_name: c_ffmpeg_decoder type:  path:  entry: \n",
                        "[2023-06-02 02:48:11.741] [info] this line: /usr/local/lib/python3.10/dist-packages/bmf/lib/libbmf_module_sdk.so.3.0\n",
                        "[2023-06-02 02:48:11.741] [info] c_ffmpeg_decoder c++ /usr/local/lib/python3.10/dist-packages/bmf/lib/libbuiltin_modules.so libbuiltin_modules.CFFDecoder\n",
                        "[2023-06-02 02:48:11.741] [info] Module info c_ffmpeg_decoder c++ libbuiltin_modules.CFFDecoder /usr/local/lib/python3.10/dist-packages/bmf/lib/libbuiltin_modules.so\n",
                        "[2023-06-02 02:48:11.746] [info] Constructing c++ module\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "Input #0, mov,mp4,m4a,3gp,3g2,mj2, from 'bbb_360_20s.mp4':\n",
                        "  Metadata:\n",
                        "    major_brand     : isom\n",
                        "    minor_version   : 512\n",
                        "    compatible_brands: isomiso2avc1mp41\n",
                        "    encoder         : Lavf59.27.100\n",
                        "  Duration: 00:00:20.13, start: 0.008005, bitrate: 449 kb/s\n",
                        "    Stream #0:0(und): Video: h264 (High) (avc1 / 0x31637661), yuv420p(tv, bt709), 640x360 [SAR 1:1 DAR 16:9], 321 kb/s, 30 fps, 30 tbr, 15360 tbn, 60 tbc (default)\n",
                        "    Metadata:\n",
                        "      handler_name    : ISO Media file produced by Google Inc. Created on: 09/04/2022.\n",
                        "      encoder         : Lavc59.37.100 libx264\n",
                        "    Stream #0:1(und): Audio: aac (LC) (mp4a / 0x6134706D), 44100 Hz, stereo, fltp, 128 kb/s (default)\n",
                        "    Metadata:\n",
                        "      handler_name    : ISO Media file produced by Google Inc. Created on: 09/04/2022.\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "[2023-06-02 02:48:11.864] [info] c++ module constructed\n",
                        "[2023-06-02 02:48:11.869] [info] debug module_name: EnhanceModule type:  path:  entry: __main__.EnhanceModule\n",
                        "[2023-06-02 02:48:11.869] [info] Module info EnhanceModule python __main__.EnhanceModule /content\n",
                        "Downloading: \"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth\" to /content/weights/realesr-animevideov3.pth\n",
                        "\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "100%|██████████| 2.39M/2.39M [00:00<00:00, 188MB/s]\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "[2023-06-02 02:48:16.759] [info] debug module_name: c_ffmpeg_filter type:  path:  entry: \n",
                        "[2023-06-02 02:48:16.759] [info] c_ffmpeg_filter c++ /usr/local/lib/python3.10/dist-packages/bmf/lib/libbuiltin_modules.so libbuiltin_modules.CFFFilter\n",
                        "[2023-06-02 02:48:16.759] [info] Module info c_ffmpeg_filter c++ libbuiltin_modules.CFFFilter /usr/local/lib/python3.10/dist-packages/bmf/lib/libbuiltin_modules.so\n",
                        "[2023-06-02 02:48:16.759] [info] Constructing c++ module\n",
                        "[2023-06-02 02:48:16.760] [info] c++ module constructed\n",
                        "[2023-06-02 02:48:16.760] [info] debug module_name: CompositionModule type:  path:  entry: __main__.CompositionModule\n",
                        "[2023-06-02 02:48:16.760] [info] Module info CompositionModule python __main__.CompositionModule /content\n",
                        "[2023-06-02 02:48:16.763] [info] debug module_name: c_ffmpeg_encoder type:  path:  entry: \n",
                        "[2023-06-02 02:48:16.763] [info] c_ffmpeg_encoder c++ /usr/local/lib/python3.10/dist-packages/bmf/lib/libbuiltin_modules.so libbuiltin_modules.CFFEncoder\n",
                        "[2023-06-02 02:48:16.763] [info] Module info c_ffmpeg_encoder c++ libbuiltin_modules.CFFEncoder /usr/local/lib/python3.10/dist-packages/bmf/lib/libbuiltin_modules.so\n",
                        "[2023-06-02 02:48:16.764] [info] Constructing c++ module\n",
                        "[2023-06-02 02:48:16.764] [info] node id:4 encode setting log level to: info\n",
                        "[2023-06-02 02:48:16.764] [info] c++ module constructed\n",
                        "[2023-06-02 02:48:16.765] [info] debug module_name: c_ffmpeg_encoder type:  path:  entry: \n",
                        "[2023-06-02 02:48:16.765] [info] Constructing c++ module\n",
                        "[2023-06-02 02:48:16.765] [info] node id:5 encode setting log level to: info\n",
                        "[2023-06-02 02:48:16.765] [info] c++ module constructed\n",
                        "[2023-06-02 02:48:16.768] [info] BMF Version: 3.0.0\n",
                        "[2023-06-02 02:48:16.768] [info] BMF Commit: \n",
                        "[2023-06-02 02:48:16.768] [info] start init graph\n",
                        "[2023-06-02 02:48:16.768] [info] scheduler count2\n",
                        "[2023-06-02 02:48:16.768] [info] node:c_ffmpeg_decoder 0 scheduler 0\n",
                        "[2023-06-02 02:48:16.774] [info] node:EnhanceModule 1 scheduler 0\n",
                        "[2023-06-02 02:48:16.779] [info] node:CompositionModule 3 scheduler 0\n",
                        "[2023-06-02 02:48:16.779] [info] node:c_ffmpeg_encoder 4 scheduler 1\n",
                        "[2023-06-02 02:48:16.780] [info] node:c_ffmpeg_encoder 5 scheduler 1\n",
                        "[2023-06-02 02:48:16.780] [info] debug module_name: c_ffmpeg_filter type:  path:  entry: \n",
                        "[2023-06-02 02:48:16.780] [info] Constructing c++ module\n",
                        "[2023-06-02 02:48:16.780] [info] c++ module constructed\n",
                        "[2023-06-02 02:48:16.780] [info] node:c_ffmpeg_filter 2 scheduler 0\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "[libx264 @ 0x7fb31026bd00] ignoring invalid SAR: 0/0\n",
                        "[libx264 @ 0x7fb31026bd00] using cpu capabilities: MMX2 SSE2Fast SSSE3 SSE4.2 AVX FMA3 BMI2 AVX2\n",
                        "[libx264 @ 0x7fb31026bd00] profile High, level 3.1\n",
                        "[libx264 @ 0x7fb31026bd00] 264 - core 155 r2917 0a84d98 - H.264/MPEG-4 AVC codec - Copyleft 2003-2018 - http://www.videolan.org/x264.html - options: cabac=1 ref=1 deblock=1:0:0 analyse=0x3:0x113 me=hex subme=2 psy=1 psy_rd=1.00:0.00 mixed_ref=0 me_range=16 chroma_me=1 trellis=0 8x8dct=1 cqm=0 deadzone=21,11 fast_pskip=1 chroma_qp_offset=0 threads=3 lookahead_threads=1 sliced_threads=0 nr=0 decimate=1 interlaced=0 bluray_compat=0 constrained_intra=0 bframes=0 weightp=1 keyint=50 keyint_min=5 scenecut=40 intra_refresh=0 rc_lookahead=10 rc=crf mbtree=1 crf=23.0 qcomp=0.60 qpmin=0 qpmax=69 qpstep=4 ip_ratio=1.40 aq=1:1.00\n",
                        "Output #0, mp4, to 'output.mp4':\n",
                        "  Metadata:\n",
                        "    encoder         : Lavf58.29.100\n",
                        "    Stream #0:0: Audio: aac (LC) (mp4a / 0x6134706D), 44100 Hz, stereo, fltp, 128 kb/s\n",
                        "    Stream #0:1: Video: h264 (avc1 / 0x31637661), yuv420p, 1280x720, q=2-31, 30 fps, 30 tbr, 15360 tbn, 30 tbc\n",
                        "[libx264 @ 0x7fb3107ffac0] ignoring invalid SAR: 0/0\n",
                        "[libx264 @ 0x7fb3107ffac0] using cpu capabilities: MMX2 SSE2Fast SSSE3 SSE4.2 AVX FMA3 BMI2 AVX2\n",
                        "[libx264 @ 0x7fb3107ffac0] profile High, level 3.1\n",
                        "[libx264 @ 0x7fb3107ffac0] 264 - core 155 r2917 0a84d98 - H.264/MPEG-4 AVC codec - Copyleft 2003-2018 - http://www.videolan.org/x264.html - options: cabac=1 ref=1 deblock=1:0:0 analyse=0x3:0x113 me=hex subme=2 psy=1 psy_rd=1.00:0.00 mixed_ref=0 me_range=16 chroma_me=1 trellis=0 8x8dct=1 cqm=0 deadzone=21,11 fast_pskip=1 chroma_qp_offset=0 threads=3 lookahead_threads=1 sliced_threads=0 nr=0 decimate=1 interlaced=0 bluray_compat=0 constrained_intra=0 bframes=0 weightp=1 keyint=50 keyint_min=5 scenecut=40 intra_refresh=0 rc_lookahead=10 rc=crf mbtree=1 crf=23.0 qcomp=0.60 qpmin=0 qpmax=69 qpstep=4 ip_ratio=1.40 aq=1:1.00\n",
                        "Output #0, mp4, to 'output_compose.mp4':\n",
                        "  Metadata:\n",
                        "    encoder         : Lavf58.29.100\n",
                        "    Stream #0:0: Audio: aac (LC) (mp4a / 0x6134706D), 44100 Hz, stereo, fltp, 128 kb/s\n",
                        "    Stream #0:1: Video: h264 (avc1 / 0x31637661), yuv420p, 1280x720, q=2-31, 30 fps, 30 tbr, 15360 tbn, 30 tbc\n",
                        "INFO:main:0.000000 -- (1) -- process num:  10\n",
                        "INFO:main:2.398019 -- (1) -- process num:  20\n",
                        "INFO:main:4.810669 -- (1) -- process num:  30\n",
                        "INFO:main:7.206777 -- (1) -- process num:  40\n",
                        "INFO:main:10.278919 -- (1) -- process num:  50\n",
                        "INFO:main:12.654365 -- (1) -- process num:  60\n",
                        "INFO:main:15.011752 -- (1) -- process num:  70\n",
                        "INFO:main:17.391072 -- (1) -- process num:  80\n",
                        "INFO:main:19.857616 -- (1) -- process num:  90\n",
                        "INFO:main:23.022989 -- (1) -- process num:  100\n",
                        "INFO:main:25.825473 -- (1) -- process num:  110\n",
                        "INFO:main:28.243883 -- (1) -- process num:  120\n",
                        "INFO:main:30.643347 -- (1) -- process num:  130\n",
                        "INFO:main:33.066066 -- (1) -- process num:  140\n",
                        "INFO:main:36.207258 -- (1) -- process num:  150\n",
                        "INFO:main:38.635260 -- (1) -- process num:  160\n",
                        "INFO:main:41.079730 -- (1) -- process num:  170\n",
                        "INFO:main:43.537011 -- (1) -- process num:  180\n",
                        "INFO:main:45.982688 -- (1) -- process num:  190\n",
                        "INFO:main:49.172305 -- (1) -- process num:  200\n",
                        "INFO:main:51.597147 -- (1) -- process num:  210\n",
                        "INFO:main:54.042892 -- (1) -- process num:  220\n",
                        "INFO:main:56.499944 -- (1) -- process num:  230\n",
                        "INFO:main:58.936040 -- (1) -- process num:  240\n",
                        "INFO:main:62.031521 -- (1) -- process num:  250\n",
                        "INFO:main:64.476270 -- (1) -- process num:  260\n",
                        "INFO:main:66.894277 -- (1) -- process num:  270\n",
                        "INFO:main:69.313925 -- (1) -- process num:  280\n",
                        "INFO:main:71.740559 -- (1) -- process num:  290\n",
                        "INFO:main:74.907544 -- (1) -- process num:  300\n",
                        "INFO:main:77.400694 -- (1) -- process num:  310\n",
                        "INFO:main:79.881896 -- (1) -- process num:  320\n",
                        "INFO:main:82.283349 -- (1) -- process num:  330\n",
                        "INFO:main:84.678407 -- (1) -- process num:  340\n",
                        "INFO:main:87.802527 -- (1) -- process num:  350\n",
                        "INFO:main:90.219833 -- (1) -- process num:  360\n",
                        "INFO:main:92.639750 -- (1) -- process num:  370\n",
                        "INFO:main:95.035831 -- (1) -- process num:  380\n",
                        "INFO:main:97.497821 -- (1) -- process num:  390\n",
                        "INFO:main:100.638093 -- (1) -- process num:  400\n",
                        "INFO:main:103.100864 -- (1) -- process num:  410\n",
                        "INFO:main:105.527182 -- (1) -- process num:  420\n",
                        "INFO:main:107.974263 -- (1) -- process num:  430\n",
                        "INFO:main:110.483721 -- (1) -- process num:  440\n",
                        "INFO:main:113.675591 -- (1) -- process num:  450\n",
                        "INFO:main:116.198576 -- (1) -- process num:  460\n",
                        "INFO:main:118.638803 -- (1) -- process num:  470\n",
                        "INFO:main:121.104660 -- (1) -- process num:  480\n",
                        "INFO:main:123.574502 -- (1) -- process num:  490\n",
                        "INFO:main:126.723809 -- (1) -- process num:  500\n",
                        "INFO:main:129.189348 -- (1) -- process num:  510\n",
                        "INFO:main:131.665654 -- (1) -- process num:  520\n",
                        "INFO:main:134.138080 -- (1) -- process num:  530\n",
                        "INFO:main:136.613679 -- (1) -- process num:  540\n",
                        "INFO:main:139.850631 -- (1) -- process num:  550\n",
                        "INFO:main:142.344423 -- (1) -- process num:  560\n",
                        "INFO:main:144.785160 -- (1) -- process num:  570\n",
                        "INFO:main:147.273776 -- (1) -- process num:  580\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "[2023-06-02 02:50:54.633] [info] node id:0 decode flushing\n",
                        "[2023-06-02 02:50:54.633] [info] node id:0 Process node end\n",
                        "[2023-06-02 02:50:54.634] [info] node id:0 close node\n",
                        "[2023-06-02 02:50:54.634] [info] node 0 close report, closed count: 1\n",
                        "[2023-06-02 02:50:54.634] [info] node id:1 eof received\n",
                        "[2023-06-02 02:50:54.634] [info] node id:1 eof processed, remove node from scheduler\n",
                        "[2023-06-02 02:50:54.634] [info] node id:2 eof received\n",
                        "[2023-06-02 02:50:54.634] [info] node id:2 eof processed, remove node from scheduler\n",
                        "[2023-06-02 02:50:54.634] [info] node id:4 eof received\n",
                        "[2023-06-02 02:50:54.634] [info] node id:4 eof processed, remove node from scheduler\n",
                        "[2023-06-02 02:50:54.634] [info] node id:5 eof received\n",
                        "[2023-06-02 02:50:54.634] [info] node id:5 eof processed, remove node from scheduler\n",
                        "[2023-06-02 02:50:54.634] [info] node id:1 process eof, add node to scheduler\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "INFO:main:149.648973 -- (1) -- process num:  590\n",
                        "INFO:main:149.650781 -- (1) -- Receive EOF\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "[2023-06-02 02:50:55.535] [info] node id:1 Process node end\n",
                        "[2023-06-02 02:50:55.535] [info] node id:1 close node\n",
                        "[2023-06-02 02:50:55.535] [info] node 1 close report, closed count: 2\n",
                        "[2023-06-02 02:50:55.535] [info] node id:3 eof received\n",
                        "[2023-06-02 02:50:55.536] [info] node id:4 eof received\n",
                        "[2023-06-02 02:50:55.536] [info] node id:4 eof processed, remove node from scheduler\n",
                        "[2023-06-02 02:50:55.536] [info] node id:4 process eof, add node to scheduler\n",
                        "[2023-06-02 02:50:55.536] [info] node id:2 process eof, add node to scheduler\n",
                        "[2023-06-02 02:50:55.543] [info] node id:2 Process node end\n",
                        "[2023-06-02 02:50:55.543] [info] node id:2 close node\n",
                        "[2023-06-02 02:50:55.543] [info] node 2 close report, closed count: 3\n",
                        "[2023-06-02 02:50:55.543] [info] node id:3 eof received\n",
                        "[2023-06-02 02:50:55.583] [info] node id:3 eof processed, remove node from scheduler\n",
                        "[2023-06-02 02:50:55.583] [info] node id:3 eof processed, remove node from scheduler\n",
                        "[2023-06-02 02:50:55.622] [info] node id:3 process eof, add node to scheduler\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "INFO:main:149.746040 -- (3) -- Receive EOF.................\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "[2023-06-02 02:50:55.635] [info] node id:3 Process node end\n",
                        "[2023-06-02 02:50:55.635] [info] node id:3 close node\n",
                        "[2023-06-02 02:50:55.635] [info] node 3 close report, closed count: 4\n",
                        "[2023-06-02 02:50:55.635] [info] node id:5 eof received\n",
                        "[2023-06-02 02:50:55.635] [info] node id:5 eof processed, remove node from scheduler\n",
                        "[2023-06-02 02:50:55.769] [info] node id:4 Process node end\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "[libx264 @ 0x7fb31026bd00] frame I:14    Avg QP:18.71  size: 59082\n",
                        "[libx264 @ 0x7fb31026bd00] frame P:576   Avg QP:22.48  size:  2176\n",
                        "[libx264 @ 0x7fb31026bd00] mb I  I16..4: 17.2% 41.3% 41.5%\n",
                        "[libx264 @ 0x7fb31026bd00] mb P  I16..4:  0.9%  1.9%  0.2%  P16..4:  6.2%  2.7%  1.1%  0.0%  0.0%    skip:86.9%\n",
                        "[libx264 @ 0x7fb31026bd00] 8x8 transform intra:53.4% inter:34.7%\n",
                        "[libx264 @ 0x7fb31026bd00] coded y,uvDC,uvAC intra: 52.2% 57.3% 25.5% inter: 1.7% 2.6% 0.1%\n",
                        "[libx264 @ 0x7fb31026bd00] i16 v,h,dc,p: 42% 25%  8% 25%\n",
                        "[libx264 @ 0x7fb31026bd00] i8 v,h,dc,ddl,ddr,vr,hd,vl,hu: 20% 22% 22%  6%  6%  7%  6%  6%  6%\n",
                        "[libx264 @ 0x7fb31026bd00] i4 v,h,dc,ddl,ddr,vr,hd,vl,hu: 29% 15%  9%  6%  8% 10%  7%  9%  6%\n",
                        "[libx264 @ 0x7fb31026bd00] i8c dc,h,v,p: 43% 31% 18%  9%\n",
                        "[libx264 @ 0x7fb31026bd00] Weighted P-Frames: Y:0.0% UV:0.0%\n",
                        "[libx264 @ 0x7fb31026bd00] kb/s:841.92\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "[2023-06-02 02:50:55.771] [info] node id:4 close node\n",
                        "[2023-06-02 02:50:55.771] [info] node 4 close report, closed count: 5\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "[aac @ 0x7fb3100021c0] Qavg: 255.798\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "[2023-06-02 02:50:55.839] [info] node id:5 process eof, add node to scheduler\n",
                        "[2023-06-02 02:50:56.013] [info] node id:5 Process node end\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "[libx264 @ 0x7fb3107ffac0] frame I:14    Avg QP:18.64  size: 55689\n",
                        "[libx264 @ 0x7fb3107ffac0] frame P:576   Avg QP:24.60  size:  2620\n",
                        "[libx264 @ 0x7fb3107ffac0] mb I  I16..4: 18.0% 40.1% 41.9%\n",
                        "[libx264 @ 0x7fb3107ffac0] mb P  I16..4:  1.4%  2.1%  0.3%  P16..4:  6.9%  3.3%  1.3%  0.0%  0.0%    skip:84.7%\n",
                        "[libx264 @ 0x7fb3107ffac0] 8x8 transform intra:49.9% inter:37.3%\n",
                        "[libx264 @ 0x7fb3107ffac0] coded y,uvDC,uvAC intra: 52.8% 60.3% 24.7% inter: 1.9% 3.2% 0.3%\n",
                        "[libx264 @ 0x7fb3107ffac0] i16 v,h,dc,p: 54% 23%  9% 14%\n",
                        "[libx264 @ 0x7fb3107ffac0] i8 v,h,dc,ddl,ddr,vr,hd,vl,hu: 21% 23% 19%  5%  6%  7%  6%  7%  6%\n",
                        "[libx264 @ 0x7fb3107ffac0] i4 v,h,dc,ddl,ddr,vr,hd,vl,hu: 31% 16%  9%  5%  8%  9%  7%  9%  6%\n",
                        "[libx264 @ 0x7fb3107ffac0] i8c dc,h,v,p: 38% 27% 27%  8%\n",
                        "[libx264 @ 0x7fb3107ffac0] Weighted P-Frames: Y:0.0% UV:0.0%\n",
                        "[libx264 @ 0x7fb3107ffac0] kb/s:926.31\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "[2023-06-02 02:50:56.016] [info] node id:5 close node\n",
                        "[2023-06-02 02:50:56.016] [info] node 5 close report, closed count: 6\n",
                        "[2023-06-02 02:50:56.016] [info] schedule queue 0 start to join thread\n",
                        "[2023-06-02 02:50:56.016] [info] schedule queue 0 thread quit\n",
                        "[2023-06-02 02:50:56.017] [info] schedule queue 0 closed\n",
                        "[2023-06-02 02:50:56.017] [info] schedule queue 1 start to join thread\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "[aac @ 0x7fb31012dc80] Qavg: 255.798\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "[2023-06-02 02:50:56.019] [info] schedule queue 1 thread quit\n",
                        "[2023-06-02 02:50:56.020] [info] schedule queue 1 closed\n",
                        "[2023-06-02 02:50:56.020] [info] all scheduling threads were joint\n"
                    ]
                }
            ],
            "source": [
                "#!/usr/bin/env python\n",
                "# -*- coding: utf-8 -*-\n",
                "\n",
                "input_file = \"bbb_360_20s.mp4\"\n",
                "output_file = \"output.mp4\"\n",
                "output_compose_file = \"output_compose.mp4\"\n",
                "\n",
                "import bmf\n",
                "\n",
                "\n",
                "def run():\n",
                "    graph = bmf.graph()\n",
                "    video = graph.decode(\n",
                "        {\n",
                "            \"input_path\": input_file,\n",
                "        }\n",
                "    )\n",
                "\n",
                "    enhance = bmf.module(\n",
                "        [video[\"video\"]],\n",
                "        \"EnhanceModule\",\n",
                "        option={\"fp32\": True, \"output_scale\": 2},\n",
                "        entry=\"__main__.EnhanceModule\",\n",
                "        input_manager=\"immediate\",\n",
                "    )\n",
                "\n",
                "    origin_video_scale = video[\"video\"].scale(\"2*in_w:2*in_h\")\n",
                "\n",
                "    com = bmf.module(\n",
                "        [enhance[0], origin_video_scale],\n",
                "        \"CompositionModule\",\n",
                "        entry=\"__main__.CompositionModule\",\n",
                "        input_manager=\"framesync\",\n",
                "    )\n",
                "\n",
                "    bmf.encode(\n",
                "        enhance[0],\n",
                "        video[\"audio\"],\n",
                "        {\n",
                "            \"video_params\": {\"g\": \"50\", \"preset\": \"veryfast\", \"bf\": 0, \"vsync\":\"vfr\", \"max_fr\":30},\n",
                "            \"audio_params\": {\"sample_rate\": 44100, \"codec\": \"aac\"},\n",
                "            \"loglevel\": \"info\",\n",
                "            \"output_path\": output_file,\n",
                "        },\n",
                "    )\n",
                "\n",
                "    bmf.encode(\n",
                "        com[0],\n",
                "        video[\"audio\"],\n",
                "        {\n",
                "            \"video_params\": {\"g\": \"50\", \"preset\": \"veryfast\", \"bf\": 0, \"vsync\":\"vfr\", \"max_fr\":30},\n",
                "            \"audio_params\": {\"sample_rate\": 44100, \"codec\": \"aac\"},\n",
                "            \"loglevel\": \"info\",\n",
                "            \"output_path\": output_compose_file,\n",
                "        },\n",
                "    )\n",
                "    graph.run()\n",
                "\n",
                "\n",
                "if __name__ == \"__main__\":\n",
                "    run()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "GkgriFBYG7uV"
            },
            "source": [
                "# 3 Visulize output video"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 471,
                    "output_embedded_package_id": "1o3Ho8eI4u3jnHJXsq5lS-YVcVrnSO88m"
                },
                "executionInfo": {
                    "elapsed": 3500,
                    "status": "ok",
                    "timestamp": 1685522738622,
                    "user": {
                        "displayName": "hu heng",
                        "userId": "04066773367750540207"
                    },
                    "user_tz": -480
                },
                "id": "RbHzvolvG9ae",
                "outputId": "be37175a-38a3-4e37-c0d0-61cf107dad59"
            },
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "Output hidden; open in https://colab.research.google.com to view."
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                }
            ],
            "source": [
                "from IPython.display import HTML\n",
                "from base64 import b64encode\n",
                " \n",
                "def show_video(video_path, video_width = 800):\n",
                "   \n",
                "  video_file = open(video_path, \"r+b\").read()\n",
                " \n",
                "  video_url = f\"data:video/mp4;base64,{b64encode(video_file).decode()}\"\n",
                "  return HTML(f\"\"\"<video width={video_width} controls><source src=\"{video_url}\"></video>\"\"\")\n",
                "\n",
                "# input video\n",
                "show_video('output_compose.mp4')"
            ]
        }
    ],
    "metadata": {
        "accelerator": "GPU",
        "colab": {
            "authorship_tag": "ABX9TyOt8Bz0Cw/HpyviYHfRCx/P",
            "collapsed_sections": [
                "6NMSi4HFFtJk"
            ],
            "provenance": [
                {
                    "file_id": "1CzYWlqtQGWnt6CsXhEbPLFegckEG0A2t",
                    "timestamp": 1685511238504
                }
            ]
        },
        "gpuClass": "standard",
        "kernelspec": {
            "display_name": "Python 3",
            "name": "python3"
        },
        "language_info": {
            "name": "python"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 0
}
